"""
思维链生成服务
"""
import logging
import json
from typing import List, Dict, Any
from uuid import uuid4

from backend.services.ai_service import get_ai_service
from backend.models.core import Task, QuestionAnalysis, ValidationResult
from shared.enums import TaskType, AICapability, TaskStatus

logger = logging.getLogger(__name__)


class ChainOfThoughtGenerator:
    """思维链任务生成器"""
    
    def __init__(self):
        self.task_generation_prompt = """
基于以下问题和分析结果，生成一个结构化的思维链任务序列。

问题：{question}

问题分析：
- 复杂度：{complexity}
- 类型：{category}
- 关键主题：{key_topics}
- 建议方式：{suggested_approach}

请生成3-6个具体的执行任务，每个任务应该：
1. 有明确的描述和目标
2. 指定任务类型
3. 定义所需的AI能力
4. 设置依赖关系（如果有）
5. 预估执行时间（秒）
6. 设置优先级（1-5，数字越小优先级越高）

任务类型包括：
- research: 研究任务
- analysis: 分析任务
- synthesis: 综合任务
- generation: 生成任务
- evaluation: 评估任务
- comparison: 比较任务

AI能力包括：
- text_generation: 文本生成
- analysis: 分析能力
- reasoning: 推理能力
- research: 研究能力
- creativity: 创造力
- translation: 翻译能力

请返回JSON格式的任务列表：
{{
    "tasks": [
        {{
            "description": "任务描述",
            "type": "任务类型",
            "required_capabilities": ["能力1", "能力2"],
            "dependencies": ["依赖任务的ID"],
            "estimated_duration": 数字（秒）,
            "priority": 数字（1-5）
        }}
    ]
}}

注意：
- 第一个任务通常没有依赖
- 后续任务可以依赖前面的任务
- 依赖关系用任务在数组中的索引表示（0, 1, 2...）
- 确保任务序列逻辑合理

只返回JSON，不要其他内容。
"""
    
    async def generate_tasks(self, question: str, analysis: QuestionAnalysis) -> List[Task]:
        """生成任务序列"""
        try:
            prompt = self.task_generation_prompt.format(
                question=question,
                complexity=analysis.complexity.value,
                category=analysis.category.value,
                key_topics=", ".join(analysis.key_topics),
                suggested_approach=analysis.suggested_approach
            )
            
            result = await get_ai_service().generate_text(
                prompt,
                temperature=0.4,
                max_tokens=1500
            )
            
            if not result.success:
                logger.error(f"Task generation failed: {result.error}")
                return self._get_default_tasks(question)
            
            # 解析JSON响应
            try:
                task_data = json.loads(result.result.strip())
                tasks = self._parse_task_list(task_data.get('tasks', []))
                
                # 优化任务序列
                optimized_tasks = await self.optimize_task_sequence(tasks)
                
                logger.info(f"Generated {len(optimized_tasks)} tasks")
                return optimized_tasks
                
            except json.JSONDecodeError as e:
                logger.error(f"Failed to parse task JSON: {e}")
                logger.error(f"AI response: {result.result}")
                return self._get_default_tasks(question)
                
        except Exception as e:
            logger.error(f"Task generation error: {e}")
            return self._get_default_tasks(question)
    
    def _parse_task_list(self, task_data_list: List[Dict[str, Any]]) -> List[Task]:
        """解析任务列表"""
        tasks = []
        task_id_map = {}  # 索引到ID的映射
        
        # 第一遍：创建所有任务并建立ID映射
        for i, task_data in enumerate(task_data_list):
            task_id = str(uuid4())
            task_id_map[i] = task_id
            
            task = Task(
                id=task_id,
                description=task_data.get('description', f'任务 {i+1}'),
                type=self._parse_task_type(task_data.get('type', 'research')),
                required_capabilities=self._parse_capabilities(task_data.get('required_capabilities', [])),
                estimated_duration=task_data.get('estimated_duration', 60),
                priority=task_data.get('priority', 3),
                status=TaskStatus.PENDING,
                dependencies=[]  # 稍后设置
            )
            tasks.append(task)
        
        # 第二遍：设置依赖关系
        for i, task in enumerate(tasks):
            task_data = task_data_list[i]
            dependencies = task_data.get('dependencies', [])
            
            for dep_index in dependencies:
                if isinstance(dep_index, int) and 0 <= dep_index < len(tasks):
                    task.dependencies.append(task_id_map[dep_index])
        
        return tasks
    
    def _parse_task_type(self, type_str: str) -> TaskType:
        """解析任务类型"""
        try:
            return TaskType(type_str.lower())
        except ValueError:
            return TaskType.RESEARCH
    
    def _parse_capabilities(self, cap_list: List[str]) -> List[AICapability]:
        """解析AI能力列表"""
        capabilities = []
        for cap_str in cap_list:
            try:
                capabilities.append(AICapability(cap_str.lower()))
            except ValueError:
                continue
        
        if not capabilities:
            capabilities.append(AICapability.TEXT_GENERATION)
        
        return capabilities
    
    async def optimize_task_sequence(self, tasks: List[Task]) -> List[Task]:
        """优化任务序列"""
        # 按依赖关系和优先级排序
        def sort_key(task: Task) -> tuple:
            # 没有依赖的任务优先
            dep_count = len(task.dependencies)
            return (dep_count, task.priority, task.estimated_duration)
        
        # 验证依赖关系
        validation = self._validate_dependencies(tasks)
        if not validation.is_valid:
            logger.warning(f"Task dependencies validation failed: {validation.errors}")
            # 移除无效依赖
            self._fix_dependencies(tasks, validation)
        
        # 排序任务
        sorted_tasks = sorted(tasks, key=sort_key)
        
        logger.debug(f"Optimized task sequence: {[task.description[:30] for task in sorted_tasks]}")
        return sorted_tasks
    
    def _validate_dependencies(self, tasks: List[Task]) -> ValidationResult:
        """验证任务依赖关系"""
        result = ValidationResult(is_valid=True)
        task_ids = {task.id for task in tasks}
        
        for task in tasks:
            # 检查依赖的任务是否存在
            for dep_id in task.dependencies:
                if dep_id not in task_ids:
                    result.add_error(f"Task {task.id} depends on non-existent task {dep_id}")
            
            # 检查循环依赖（简单检查）
            if task.id in task.dependencies:
                result.add_error(f"Task {task.id} has self-dependency")
        
        return result
    
    def _fix_dependencies(self, tasks: List[Task], validation: ValidationResult):
        """修复依赖关系问题"""
        task_ids = {task.id for task in tasks}
        
        for task in tasks:
            # 移除无效依赖
            valid_deps = [dep_id for dep_id in task.dependencies if dep_id in task_ids and dep_id != task.id]
            task.dependencies = valid_deps
    
    def _get_default_tasks(self, question: str) -> List[Task]:
        """获取默认任务序列"""
        return [
            Task(
                description=f"分析问题：{question[:50]}...",
                type=TaskType.ANALYSIS,
                required_capabilities=[AICapability.ANALYSIS, AICapability.REASONING],
                estimated_duration=60,
                priority=1
            ),
            Task(
                description="收集相关信息和资料",
                type=TaskType.RESEARCH,
                required_capabilities=[AICapability.RESEARCH, AICapability.TEXT_GENERATION],
                estimated_duration=90,
                priority=2,
                dependencies=[]
            ),
            Task(
                description="综合分析并生成答案",
                type=TaskType.SYNTHESIS,
                required_capabilities=[AICapability.TEXT_GENERATION, AICapability.REASONING],
                estimated_duration=120,
                priority=3,
                dependencies=[]
            )
        ]


# 全局实例
cot_generator = ChainOfThoughtGenerator()